# -*- coding: utf-8 -*-
# @Time : 2020/10/15 09:51
# @Author : wudeyang
# @email :wudeyang@sjtu.edu.cn
# @Description: 制作数据标签

import os

'''
generate the txt file for train/test in our framework
'''


############ Utility functions ############
def get_images(img_path):
    """
    find image files in data path
    :return: list of files found
    """
    img_path = os.path.abspath(img_path)
    files = []
    exts = ['jpg', 'png', 'jpeg', 'JPG', 'PNG']
    for parent, dirnames, filenames in os.walk(img_path):
        for filename in filenames:
            for ext in exts:
                if filename.endswith(ext):
                    files.append(os.path.join(parent, filename))
                    break
    print('Find {} images'.format(len(files)))
    return sorted(files)


def get_txts(txt_path):
    """
    find gt files in data path
    :return: list of files found
    """
    txt_path = os.path.abspath(txt_path)
    files = []
    exts = ['txt']
    for parent, dirnames, filenames in os.walk(txt_path):
        for filename in filenames:
            for ext in exts:
                if filename.endswith(ext):
                    files.append(os.path.join(parent, filename))
                    break
    print('Find {} txts'.format(len(files)))
    return sorted(files)


if __name__ == '__main__':
    save_path = '/data/ICDAR_2015/dataset_ic15_test.txt'
    img_path = '/data/ICDAR_2015/test_img'

    files = get_images(img_path)
    txt_path = '/data/ICDAR_2015/test_gt'
    txts = get_txts(txt_path)
    n = len(files)
    assert len(files) == len(txts)

    with open(save_path, 'w') as f:
        for i in range(n):
            line = files[i] + '\t' + txts[i] + '\n'
            f.write(line)

    print('dataset generated ^_^ ')
